module Environment where

import Control.Monad.State
import Absfgrammar
import List
import System.Exit(exitFailure)

type Sig = [(Ident, Def)]
type TState a = StateT Environment IO a
type Selector a = (TState a, a -> TState ())

data Environment = Environment {
    envSig :: Sig,
    transformedProgram :: Program,
    argName :: Ident,
    replacementExp :: Exp,
    replaceTerms :: [Exp]
} deriving (Eq, Ord, Show)

-- ===========================================================================
-- SELECTORS
-- ===========================================================================

signature :: Selector Sig
signature = (gets envSig, 
             \x -> modify (\vs -> vs {envSig = x}))

tProg :: Selector Program
tProg = (gets transformedProgram,
         \x -> modify (\vs -> vs {transformedProgram = x}))

argumentName :: Selector Ident
argumentName = (gets argName,
                \x -> modify (\vs -> vs {argName = x}))

repExp :: Selector Exp
repExp = (gets replacementExp,
          \x -> modify (\vs -> vs {replacementExp = x}))

repTerms :: Selector [Exp]
repTerms = (gets replaceTerms,
            \x -> modify (\vs -> vs {replaceTerms = x}))

-- ===========================================================================
-- GETTERS AND SETTERS
-- ===========================================================================

addFun :: Def -> TState ()
addFun d = case d of
    (Funr ids exp) -> do
        sig <- get' signature
        put' signature (\x -> (((head ids),d):sig))
    (Func ids exp) -> do
        sig <- get' signature
        put' signature (\x -> (((head ids),d):sig))
    _ -> return ()
        

getFun :: Ident -> TState Def
getFun id = do
    sig <- get' signature
    case (lookup id sig) of
        Just d  -> return d
        Nothing -> do 
            liftIO $ putStrLn $ "Function " ++ (show id) ++ "not decl."
            liftIO $ exitFailure

getSig :: TState [(Ident, Def)]
getSig = do
    sig <- get' signature
    return sig

getProgram :: TState Program
getProgram = do
    prog <- get' tProg
    return prog

addDef :: Def -> TState ()
addDef def = do
    (Prog ds) <- get' tProg
    put' tProg (\x -> (Prog (ds++[def])))

-- Put the function name to get replaced argument
putArgName :: Ident -> TState ()
putArgName id = put' argumentName (\x -> id)

-- Get the function name to get replaced argument
getArgName :: TState Ident
getArgName = do
    id <- get' argumentName
    return id

putReplacementExp :: Exp -> TState ()
putReplacementExp e = put' repExp (\x -> e)

getReplacementExp :: TState Exp
getReplacementExp = do
    e <- get' repExp
    return e

putReplaceTerms :: [Exp] -> TState ()
putReplaceTerms ts = put' repTerms (\x -> ts)

getReplaceTerms :: TState [Exp]
getReplaceTerms = do
    ts <- get' repTerms
    return ts

-- ===========================================================================
-- MONAD GET AND PUT
-- ===========================================================================

get' :: Selector a -> TState a
get' = fst

put' :: Selector a -> (a -> a) -> TState ()
put' (gf,uf) f = do
    st <- gf
    uf (f st)

-- ===========================================================================
-- INITIAL STATE
-- ===========================================================================

emptySignature :: Sig
emptySignature = []

emptyProgram :: Program
emptyProgram = (Prog [])

s0 :: Environment
s0 = Environment {
       envSig = emptySignature,
       transformedProgram = emptyProgram,
       argName = (Ident ""),
       replacementExp = (EId (Ident "")),
       replaceTerms = []
}
